"""
Image utilities.
"""
import io
from enum import auto, Enum

import numpy as np
import skimage.io  # type: ignore
from PIL import Image  # type: ignore

from ._typing import NDFloat32Array


class Preprocessing(Enum):
    YAHOO = auto()
    SIMPLE = auto()


def preprocess_image(
        pil_image: Image,  # 输入的PIL图像对象
        preprocessing: Preprocessing = Preprocessing.YAHOO  # 预处理策略，默认为Yahoo
) -> NDFloat32Array:
    """
    预处理用于预训练的Open NSFW模型权重的图像。

    参考:
    https://github.com/mdietrichstein/tensorflow-open_nsfw
    """
    # 确保图像是RGB格式
    if pil_image.mode != "RGB":
        pil_image = pil_image.convert("RGB")

    # 使用Yahoo的预处理方式
    if preprocessing == Preprocessing.YAHOO:
        # 将图像大小调整为256x256，使用双线性插值
        pil_image_resized = pil_image.resize(
            (256, 256), resample=Image.BILINEAR
        )

        # 创建一个内存中的字节流
        fh_im = io.BytesIO()
        # 将调整大小的图像保存到字节流中，格式为JPEG
        pil_image_resized.save(fh_im, format="JPEG")
        fh_im.seek(0)

        # 从字节流中读取图像，不转为灰度图，将图像转换为float32类型的数组
        image: NDFloat32Array = skimage.io.imread(fh_im, as_gray=False).astype(np.float32)

        # 获取图像的高度和宽度
        height, width, _ = image.shape
        # 设置最终图像的高度和宽度
        h, w = (224, 224)

        # 计算高度方向上的裁剪偏移
        h_off = max((height - h) // 2, 0)
        # 计算宽度方向上的裁剪偏移
        w_off = max((width - w) // 2, 0)
        # 裁剪图像
        image = image[h_off:h_off + h, w_off:w_off + w, :]

    # 使用简单的预处理方式
    elif preprocessing == Preprocessing.SIMPLE:
        # 将图像大小调整为224x224，使用双线性插值
        pil_image_resized = pil_image.resize(
            (224, 224), resample=Image.BILINEAR
        )
        # 将PIL图像转换为float32类型的数组
        image = np.array(pil_image_resized).astype(np.float32)

    else:
        # 如果预处理选项无效，抛出异常
        raise ValueError("Invalid preprocessing option.")

    # 将图像从RGB格式转换为BGR格式
    image = image[:, :, ::-1]

    # VGG模型的三个通道的平均值
    vgg_mean = [104, 117, 123]
    # 从每个通道中减去平均值
    image = image - np.array(vgg_mean, dtype=np.float32)

    # 返回预处理后的图像数组
    return image
